package sketch.ounit.fuzz;

import java.util.LinkedHashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

import org.eclipse.jdt.core.dom.ASTNode;
import org.eclipse.jdt.core.dom.CompilationUnit;
import org.eclipse.jdt.core.dom.DoStatement;
import org.eclipse.jdt.core.dom.EnhancedForStatement;
import org.eclipse.jdt.core.dom.Expression;
import org.eclipse.jdt.core.dom.ForStatement;
import org.eclipse.jdt.core.dom.MethodDeclaration;
import org.eclipse.jdt.core.dom.MethodInvocation;
import org.eclipse.jdt.core.dom.PackageDeclaration;
import org.eclipse.jdt.core.dom.WhileStatement;
import org.eclipse.jdt.internal.core.dom.NaiveASTFlattener;

import sketch.ast.ASTUtils;
import sketch.ounit.Values;
import sketch.util.Checker;

public class RandomFuzzExtractor {

	private final String file_path;
	private final CompilationUnit unit;
	
	private final Map<MethodDeclaration, List<Expression>> expr_map
	  = new LinkedHashMap<MethodDeclaration, List<Expression>>();
	
	private final Map<MethodDeclaration, List<MethodInvocation>> invoke_map
	  = new LinkedHashMap<MethodDeclaration, List<MethodInvocation>>();
	
	private final Map<MethodDeclaration, List<Expression[]>> obj_args_map
	  = new LinkedHashMap<MethodDeclaration, List<Expression[]>>();
	
	public RandomFuzzExtractor(String file_path) {
		Checker.checkNull(file_path, "File path could not be null.");
		this.file_path = file_path;
		this.unit = ASTUtils.astRoot(this.file_path);
		//check there must be only 1 type declaration
		Checker.checkTrue(this.unit.types().size() == 1, "The size of types should be 1, but now it is: "
				+ this.unit.types().size());
		//do the analysis
		extract_fuzzing_places();
	}
	
	private void extract_fuzzing_places() {
		List<MethodDeclaration> methods = ASTUtils.getAllMethods(unit);
		for(MethodDeclaration method : methods) {
			RandomFuzzVisitor visitor = new RandomFuzzVisitor();
			method.accept(visitor);
			if(!visitor.getReceiverExpressions().isEmpty()) {
				expr_map.put(method, visitor.getReceiverExpressions());
				invoke_map.put(method, visitor.getMethodInvocations());
				obj_args_map.put(method, visitor.getObjectArgExpressions());
			}
		}
	}
	
	public CompilationUnit getUnit() {
		return this.unit;
	}
	
	public String getPackageName() {
		PackageDeclaration packageDeclaration = this.unit.getPackage();
		if(packageDeclaration == null) {
			return "";
		} else {
			return packageDeclaration.getName().toString();
		}
	}
	
	public Map<MethodDeclaration, List<Expression>> getExprMap() {
		return this.expr_map;
	}
	
	public Map<MethodDeclaration, List<MethodInvocation>> getInvokeMap() {
		return this.invoke_map;
	}
	
	public Map<MethodDeclaration, List<Expression[]>> getArgExprMap() {
		return this.obj_args_map;
	}
	
	public static void main(String[] args) {
		String file_path = "./tests/sketch/ounit/fuzz/ExistingObjectFuzzerTest.java";
		CompilationUnit unit = ASTUtils.astRoot(file_path);
		RandomFuzzVisitor visitor = new RandomFuzzVisitor();
		visitor.check_valid = false;
		unit.accept(visitor);
		System.out.println(visitor.getReceiverExpressions());
		System.out.println(visitor.getMethodInvocations());
		System.out.println(visitor.getObjectArgExpressions());
	}
}

//the underlying visitor
class RandomFuzzVisitor extends NaiveASTFlattener {
	
	boolean check_valid = true;
	
	//the list must be in order of the receiver names as they appear
	//in the code
	private List<Expression> receiver_expressions = new LinkedList<Expression>();
	private List<MethodInvocation> method_invocations = new LinkedList<MethodInvocation>();
	private List<Expression[]> object_arg_expressions = new LinkedList<Expression[]>();
	public List<Expression> getReceiverExpressions() {
		return this.receiver_expressions;
	}
	public List<MethodInvocation> getMethodInvocations() {
		return this.method_invocations;
	}
	public List<Expression[]> getObjectArgExpressions() {
		return this.object_arg_expressions;
	}
	
	private boolean isRandomValue(String invokeExprName, String methodName) {
		return (invokeExprName.equals(Values.random_fuzz_class) || invokeExprName.equals(Values.random_fuzz_full_class))
				&& (methodName.equals("random") || methodName.equals("randomMutate") || methodName.equals("randomCreate"));
	}
	public boolean visit(MethodInvocation invoke) {
		String invokeExpression = invoke.getExpression() == null ? "" : invoke.getExpression().toString();
		String methodName = invoke.getName().toString(); 
		if(this.isRandomValue(invokeExpression, methodName)) {
			//System.out.println(invoke);
			Checker.checkTrue(!invoke.arguments().isEmpty(), "The argument could not be empty.");
			//verify the fuzz place
			if(check_valid) {
			    this.verify_fuzz_place(invoke);
			}
			if(methodName.equals("randomMutate")) {
				this.receiver_expressions.add((Expression)invoke.arguments().get(0));
				this.method_invocations.add(invoke);
				int length = invoke.arguments().size();
				  Checker.checkTrue(length >= 2, "The length of arguments should >= 2, but now it is: " + invoke);
				  Expression[] arg_expressions = new Expression[invoke.arguments().size() - 2];
				  for(int i = 2; i < invoke.arguments().size(); i++) {
					arg_expressions[i-2] = (Expression) invoke.arguments().get(i);
				  }
				  this.object_arg_expressions.add(arg_expressions);
			} else if(methodName.equals("randomCreate")) {
				this.receiver_expressions.add((Expression)invoke.arguments().get(0));
				this.method_invocations.add(invoke);
				Expression[] arg_expressions = new Expression[invoke.arguments().size() - 1];
				for(int i = 1; i < invoke.arguments().size(); i++) {
					arg_expressions[i - 1] = (Expression) invoke.arguments().get(i);
				}
				this.object_arg_expressions.add(arg_expressions);
			} else {
			  //System.out.println("   " + invoke.arguments().get(0));
			  this.receiver_expressions.add((Expression)invoke.arguments().get(0));
			  this.method_invocations.add(invoke);
			  int length = invoke.arguments().size();
			  Checker.checkTrue(length >= 2, "The length of arguments should >= 2, but now it is: " + invoke);
			  Expression[] arg_expressions = new Expression[invoke.arguments().size() - 2];
			  for(int i = 2; i < invoke.arguments().size(); i++) {
				arg_expressions[i-2] = (Expression) invoke.arguments().get(i);
			  }
			  this.object_arg_expressions.add(arg_expressions);
			}
		}
		return super.visit(invoke);
	}
	
	private void verify_fuzz_place(MethodInvocation invoke) {
		ASTNode node = invoke;
		while(node.getParent() != null) {
			node = node.getParent();
			if(node instanceof DoStatement || node instanceof ForStatement || node instanceof WhileStatement
					|| node instanceof EnhancedForStatement) {
				Checker.checkNull(null, "The method invoke could not be nested in a node: " + node.getClass());
			}
		}
	}
}